package datastructure.unionfind;

/**
 * @author MaoLin Wang
 * @description
 * @date 2020-08-24 14:47
 */
public class UF {
    private int[] parent;
    private int count;
    private int[] weight;

    public UF(int count){
        this.count=count;
        this.parent =new int[count];
        weight=new int[count];
        for (int i = 0; i < count; i++) {
            parent[i]=i;
            weight[i]=1;
        }
    }
    public boolean connected(int i ,int j){
        return pathCompressionFind(i)==pathCompressionFind(j);
    }


    public int quickFind(int i){
        if (i<0||i>count-1){
            throw new RuntimeException("Index out of Boundary");
        }
        return parent[i];
    }

    public void union(int i,int j){
        if (!connected(i,j)){
            for (int k = 0; k < parent.length; k++) {
                if (parent[k]== parent[i]){
                    parent[k]= parent[j];
                    count--;
                }
            }
        }
    }

    /**
     * id[i]保存同一个连通分量的另一个元素
     * @param i
     * @param j
     */
    public void quickUnion(int i,int j){
        int p=quickUnionOfFind(i);
        int q=quickUnionOfFind(j);
        if (p!=q){
            parent[p]=q;
            count--;
        }
    }

    private int quickUnionOfFind(int i) {
        while (i!= parent[i]){
            i= parent[i];
        }
        return i;
    }


    /**
     * 加权quickUnion,记录每个节点的高度，union时将小树合并到大树上
     */
    public void quickUnionWeight(int i,int j){
        int p = pathCompressionFind(i);
        int q = pathCompressionFind(j);
        if (p==q){
            return;
        }
        if (weight[p]<weight[q]){
            parent[p]=q;
            weight[q]+=weight[p];
        }else {
            parent[q]=p;
            weight[p]+=weight[q];
        }
        count--;
    }

    /**
     * 路径压缩，将当前节点到根节点的所有节点都改为根根节点值
     * @param i
     * @return
     */
    public int pathCompressionFind(int x){
        // 先找到根节点
        int root = x;
        while (parent[root] != root) {
            root = parent[root];
        }
        // 然后把 x 到根节点之间的所有节点直接接到根节点下面
        int old_parent = parent[x];
        while (x != root) {
            parent[x] = root;
            x = old_parent;
            old_parent = parent[old_parent];
        }
        return root;
    }
}
